#!/usr/bin/env python3
"""
Paper Search MCP Tool for Academic Paper Integration.

This module provides functionality to search academic papers using the arXiv API
and format them for use in Manim animation generation. It implements MCP-style
tool calling with clear highlighting when the tool is being used.

Usage:
    from paper_search_mcp import search_papers, format_papers_for_animation

    papers = search_papers("machine learning", max_results=3)
    context = format_papers_for_animation(papers)
"""

import arxiv
import re
import sys
from typing import List, Dict, Optional, Any
from dataclasses import dataclass
from datetime import datetime


@dataclass
class PaperInfo:
    """Data class to hold paper information."""
    title: str
    authors: List[str]
    abstract: str
    url: str
    published: datetime
    categories: List[str]
    summary_short: str


def highlight_mcp_call(message: str) -> None:
    """
    Highlight MCP tool calls with visual indicators.

    Args:
        message: The message to highlight
    """
    print("🔍 [MCP PAPER SEARCH TOOL]", message)


def clean_text(text: str) -> str:
    """
    Clean and format text for better readability.

    Args:
        text: Raw text to clean

    Returns:
        Cleaned text
    """
    # Remove extra whitespace and newlines
    text = re.sub(r'\s+', ' ', text.strip())

    # Remove LaTeX commands (basic cleanup)
    text = re.sub(r'\$[^$]*\$', '[MATH]', text)  # Inline math
    text = re.sub(r'\\\w+\{[^}]*\}', '', text)   # LaTeX commands

    return text


def extract_key_concepts(abstract: str, max_concepts: int = 5) -> List[str]:
    """
    Extract key concepts from paper abstract for animation focus.

    Args:
        abstract: Paper abstract text
        max_concepts: Maximum number of concepts to extract

    Returns:
        List of key concepts
    """
    # Simple keyword extraction - could be enhanced with NLP
    keywords = []

    # Common academic/technical terms that are good for visualization
    concept_patterns = [
        r'\b(algorithm|method|approach|technique|framework)\b',
        r'\b(model|network|system|architecture)\b',
        r'\b(optimization|learning|training|inference)\b',
        r'\b(classification|regression|clustering|prediction)\b',
        r'\b(neural|deep|machine|artificial)\b',
        r'\b(graph|tree|matrix|vector|tensor)\b',
        r'\b(probability|statistics|distribution|variance)\b',
        r'\b(gradient|derivative|function|equation)\b'
    ]

    abstract_lower = abstract.lower()
    for pattern in concept_patterns:
        matches = re.findall(pattern, abstract_lower)
        keywords.extend(matches)

    # Remove duplicates and limit
    unique_keywords = list(dict.fromkeys(keywords))[:max_concepts]
    return unique_keywords


def search_papers(query: str, max_results: int = 5, sort_by: str = "relevance") -> List[PaperInfo]:
    """
    Search for academic papers using arXiv API.

    Args:
        query: Search query string
        max_results: Maximum number of papers to return
        sort_by: Sort criteria ("relevance", "lastUpdatedDate", "submittedDate")

    Returns:
        List of PaperInfo objects
    """
    highlight_mcp_call(f"Searching for papers with query: '{query}'")

    try:
        # Configure search
        search = arxiv.Search(
            query=query,
            max_results=max_results,
            sort_by=getattr(arxiv.SortCriterion, sort_by.title(), arxiv.SortCriterion.Relevance)
        )

        papers = []

        highlight_mcp_call(f"Processing search results...")

        # Use Client.results() instead of Search.results() to avoid deprecation warning
        client = arxiv.Client()
        for result in client.results(search):
            # Create short summary (first 200 chars of abstract)
            summary_short = clean_text(result.summary)[:200] + "..." if len(result.summary) > 200 else clean_text(result.summary)

            paper = PaperInfo(
                title=clean_text(result.title),
                authors=[author.name for author in result.authors],
                abstract=clean_text(result.summary),
                url=result.entry_id,
                published=result.published,
                categories=[cat for cat in result.categories],
                summary_short=summary_short
            )

            papers.append(paper)

        highlight_mcp_call(f"Found {len(papers)} papers")
        return papers

    except Exception as e:
        highlight_mcp_call(f"Error searching papers: {str(e)}")
        return []


def format_papers_for_animation(papers: List[PaperInfo], focus_paper_index: int = 0) -> str:
    """
    Format paper information for inclusion in animation generation prompts.

    Args:
        papers: List of PaperInfo objects
        focus_paper_index: Index of the main paper to focus on (default: 0)

    Returns:
        Formatted string for use in animation prompts
    """
    if not papers:
        return ""

    highlight_mcp_call(f"Formatting {len(papers)} papers for animation context")

    # Focus on the main paper
    main_paper = papers[focus_paper_index] if focus_paper_index < len(papers) else papers[0]

    # Extract key concepts for visualization
    key_concepts = extract_key_concepts(main_paper.abstract)

    context = f"""
PAPER CONTEXT FOR ANIMATION:

Main Paper: "{main_paper.title}"
Authors: {', '.join(main_paper.authors[:3])}{'...' if len(main_paper.authors) > 3 else ''}
Published: {main_paper.published.strftime('%Y-%m-%d')}

Abstract Summary:
{main_paper.summary_short}

Key Concepts for Visualization:
{', '.join(key_concepts) if key_concepts else 'General concepts from the abstract'}

Animation Focus:
Please create an animation that visualizes the core ideas from this paper. Focus on making complex concepts accessible through visual representation. Consider using:
- Diagrams and flowcharts for algorithms/methods
- Graphs and plots for data/results
- Step-by-step animations for processes
- Mathematical visualizations for equations/formulas

Additional Context Papers:
"""

    # Add brief info about other papers
    for i, paper in enumerate(papers[1:], 1):
        if i <= 2:  # Limit to 2 additional papers
            context += f"{i}. {paper.title} ({paper.published.strftime('%Y')})\n"

    return context


def search_and_format_papers(query: str, max_results: int = 3, focus_index: int = 0) -> str:
    """
    Convenience function to search and format papers in one call.

    Args:
        query: Search query
        max_results: Maximum papers to search
        focus_index: Index of paper to focus on

    Returns:
        Formatted paper context for animation generation
    """
    highlight_mcp_call(f"Starting paper search and formatting process")

    papers = search_papers(query, max_results)
    if not papers:
        highlight_mcp_call("No papers found or search failed")
        return ""

    context = format_papers_for_animation(papers, focus_index)

    highlight_mcp_call("Paper search and formatting completed successfully")
    return context


def main():
    """Command-line interface for testing the paper search tool."""
    import argparse

    parser = argparse.ArgumentParser(description="Search academic papers for animation context")
    parser.add_argument("query", type=str, help="Search query for papers")
    parser.add_argument("--max-results", "-n", type=int, default=3, help="Maximum number of papers")
    parser.add_argument("--focus", "-f", type=int, default=0, help="Index of paper to focus on")

    args = parser.parse_args()

    context = search_and_format_papers(args.query, args.max_results, args.focus)

    if context:
        print("\n" + "="*80)
        print("GENERATED PAPER CONTEXT:")
        print("="*80)
        print(context)
    else:
        print("Failed to generate paper context.")
        return 1

    return 0


if __name__ == "__main__":
    sys.exit(main())